package com.gy.spark.sparkstreaming.offsetManager.getOffset;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryUntilElapsed;
import org.apache.kafka.common.TopicPartition;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;


public class GetTopicOffsetFromZookeeper {

    public static Map<TopicPartition, Long> getConsumerOffsets(String zkServers, String groupID, String topic) {
        Map<TopicPartition, Long> retVals = new HashMap<TopicPartition, Long>();

        ObjectMapper objectMapper = new ObjectMapper();
        CuratorFramework curatorFramework = CuratorFrameworkFactory.builder()
                .connectString(zkServers).connectionTimeoutMs(1000)
                .sessionTimeoutMs(10000).retryPolicy(new RetryUntilElapsed(1000, 1000)).build();

        curatorFramework.start();

        try {
            String nodePath = "/consumers/" + groupID + "/offsets/" + topic;
            if (curatorFramework.checkExists().forPath(nodePath) != null) {
                List<String> partitions = curatorFramework.getChildren().forPath(nodePath);
                for (String partiton : partitions) {
                    int partitionL = Integer.valueOf(partiton);
                    Long offset = objectMapper.readValue(curatorFramework.getData().forPath(nodePath + "/" + partiton), Long.class);
                    TopicPartition TopicPartition = new TopicPartition(topic, partitionL);
                    retVals.put(TopicPartition, offset);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        curatorFramework.close();

        return retVals;
    }


    public static void main(String[] args) {
        Map<TopicPartition, Long> consumerOffsets = getConsumerOffsets("node3:2181,node4:2181,node5:2181", "zhy", "mytopic");
        Set<Entry<TopicPartition, Long>> entrySet = consumerOffsets.entrySet();
        for (Entry<TopicPartition, Long> entry : entrySet) {
            TopicPartition TopicPartition = entry.getKey();
            String topic = TopicPartition.topic();
            int partition = TopicPartition.partition();
            Long offset = entry.getValue();
            System.out.println("topic = " + topic + ",partition = " + partition + ",offset = " + offset);
        }
    }
}
